import torch
from torch import Tensor

from .lib import dose_torch  # ruff: isort: skip


# 导出接口
__all__ = ["accumulate", "get_metric", "get_dose_stat_roi_interp"]

print(f'torch.ops.dose_torch.__file__ -> {dose_torch.__file__}')
print(f'dir(dose_torch) -> {dir(dose_torch)}')
print(f'dir(torch.ops) -> {dir(torch.ops)}')
print(f'dir(torch.ops.dose_torch) -> {dir(torch.ops.dose_torch)}')


##00000000000000000000000000000000000000000000000000000000000000000 提供与原始C++函数类似的接口
def accumulate(
    voxelVolume: Tensor,
    voxelDose: Tensor,
    doExactAccumulation: bool = False,
):
    """
    计算DVH(剂量体积直方图)

    Args:
        voxelVolume: 体素体积张量 (1D, double)
        voxelDose: 体素剂量张量 (1D, double)
        doExactAccumulation: 是否使用精确计算方法

    Returns:
        tuple: (binnedDoseValues, accumulatedRelativeVolumes)
    """
    # 确保输入张量类型正确
    if voxelVolume.dtype != torch.float64:
        voxelVolume = voxelVolume.double()
    if voxelDose.dtype != torch.float64:
        voxelDose = voxelDose.double()

    ## 返回 accumulatedRelativeVolumes, binnedDoseValues
    return torch.ops.dose_torch.accumulate.default(
        voxelVolume,
        voxelDose,
        doExactAccumulation,
    )


def get_metric(
    binnedDoseValues: Tensor,
    accumulatedRelativeVolumes: Tensor,
    doseAtVolume: float,
) -> Tensor:
    """
    获取特定体积百分比对应的剂量

    Args:
        doseAtVolume: 目标体积百分比 (标量, double)
        binnedDoseValues: 分箱剂量值 (1D, double)
        accumulatedRelativeVolumes: 累积相对体积 (1D, double)

    Returns:
        Tensor: 对应的剂量值
    """
    # 确保输入张量类型正确
    if False and doseAtVolume.dtype != torch.float64:
        doseAtVolume = doseAtVolume.double()
    if binnedDoseValues.dtype != torch.float64:
        binnedDoseValues = binnedDoseValues.double()
    if accumulatedRelativeVolumes.dtype != torch.float64:
        accumulatedRelativeVolumes = accumulatedRelativeVolumes.double()

    return torch.ops.dose_torch.get_metric.default(
        binnedDoseValues,
        accumulatedRelativeVolumes,
        doseAtVolume,
    )


def get_dose_stat_roi_interp(
    binnedDoseValues: Tensor,
    accumulatedRelativeVolumes: Tensor,
    voxelVolume: Tensor,
    voxelDose: Tensor,
    dose_type: str,
    doseAtVolume: float = 0.0,
) -> Tensor:
    """
    计算ROI内的剂量统计指标

    Args:
        dose_type: 剂量类型 ('average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume')
        binnedDoseValues: 分箱剂量值 (1D, double)
        accumulatedRelativeVolumes: 累积相对体积 (1D, double)
        voxelVolume: 体素体积张量 (1D, double)
        voxelDose: 体素剂量张量 (1D, double)
        doseAtVolume: 目标体积百分比 (仅当dose_type='DoseAtVolume'时使用)

    Returns:
        Tensor: 剂量统计指标值
    """
    # 确保输入张量类型正确
    if binnedDoseValues.dtype != torch.float64:
        binnedDoseValues = binnedDoseValues.double()
    if accumulatedRelativeVolumes.dtype != torch.float64:
        accumulatedRelativeVolumes = accumulatedRelativeVolumes.double()
    if voxelVolume.dtype != torch.float64:
        voxelVolume = voxelVolume.double()
    if voxelDose.dtype != torch.float64:
        voxelDose = voxelDose.double()

    return torch.ops.dose_torch.get_dose_stat_roi_interp.default(
        binnedDoseValues,
        accumulatedRelativeVolumes,
        voxelVolume,
        voxelDose,
        dose_type,
        doseAtVolume,
    )
